package com.clei.utils;

import com.clei.dto.PortUse;
import com.clei.utils.helper.TableDataPrinter;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Scanner;
import java.util.Set;
import java.util.function.Consumer;
import java.util.function.Predicate;
import java.util.stream.Collectors;

/**
 * windows 端口工具类
 *
 * @author Y
 * @date 2024-03-05
 **/
public class WinPortUtil {

    public static void main(String[] args) {
        Scanner input = new Scanner(System.in);
        while (true) {
            PrintUtil.log("请输入端口(输入0结束) : ");
            int port = input.nextInt();
            if (0 == port) {
                break;
            }
            killByPort(port);
        }
        PrintUtil.log("END");
        input.close();
    }

    /**
     * killByPort
     *
     * @param port port
     */
    public static void killByPort(Integer port) {
        List<PortUse> portUseList = list();
        Set<String> addedSet = new HashSet<>();
        List<PortUse> filterList = portUseList.stream()
                .filter(p -> port.equals(p.getPort()) && 0 != p.getPort())
                .peek(p -> addedSet.add(p.getPid()))
                .collect(Collectors.toList());
        if (CollectionUtil.isEmpty(filterList)) {
            PrintUtil.log("未查询到端口[{}]信息", port);
            return;
        }
        // 打印
        Collections.sort(filterList);
        print(filterList);

        boolean kill = true;
        boolean noAsk = false;
        Scanner input = new Scanner(System.in);
        if (addedSet.size() > 1) {
            PrintUtil.log("发现多个进程，是否全部结束(y/n)?");
            String answer = input.next();
            if (answer.equalsIgnoreCase("y")) {
                noAsk = true;
            } else {
                kill = false;
            }
        }
        if (kill) {
            Set<String> killedSet = new HashSet<>(addedSet.size());
            for (PortUse portUse : filterList) {
                if (killedSet.contains(portUse.getPid())) {
                    continue;
                }
                String process = portUse.getProcess();
                String pid = portUse.getPid();
                if (!noAsk) {
                    PrintUtil.log("是否结束进程[{}] (y/n)?", process);
                    String answer = input.next();
                    if (!answer.equalsIgnoreCase("y")) {
                        break;
                    }
                }
                CmdUtil.exec(String.format("taskkill /f /pid %s", pid));
                PrintUtil.log("已结束进程[{}]", process);
                killedSet.add(portUse.getPid());
            }
        }
    }

    /**
     * print
     */
    public static void print() {
        print(list());
    }

    /**
     * print
     */
    public static void print(Predicate<? super PortUse> filter) {
        print(list(filter));
    }

    /**
     * list
     *
     * @return list
     */
    public static List<PortUse> list() {
        return list(null);
    }

    /**
     * list
     *
     * @param filter filter
     * @return list
     */
    public static List<PortUse> list(Predicate<? super PortUse> filter) {
        // 端口信息收集
        List<PortUse> portUseList = new ArrayList<>();
        Consumer<String> portHandler = s -> {
            String[] arr = s.trim().split("\\s+");
            if (arr.length < 2 || "0".equals(arr[arr.length - 1])) {
                return;
            }
            String ipPort = arr[1];
            int index = ipPort.lastIndexOf(':');
            if (-1 != index) {
                String protocol = arr[0];
                String pid = arr[arr.length - 1];
                String state = "UDP".equals(protocol) ? "" : arr[3];
                String port = ipPort.substring(index + 1);
                portUseList.add(new PortUse(protocol, ipPort.substring(0, index), Integer.valueOf(port), state, pid));
            }
        };
        CmdUtil.exec("netstat -ano", portHandler);
        // pid对应process
        Map<String, String> pidProcessMap = new HashMap<>();
        Consumer<String> processHandler = s -> {
            String[] arr = s.trim().split("\\s+");
            if (arr.length > 1) {
                pidProcessMap.put(arr[1], arr[0]);
            }
        };
        CmdUtil.exec("tasklist", processHandler);
        // 设置对应
        for (PortUse p : portUseList) {
            p.setProcess(pidProcessMap.get(p.getPid()));
        }
        List<PortUse> resList = portUseList;
        // filter
        if (null != filter) {
            resList = portUseList.stream().filter(filter).collect(Collectors.toList());
        }
        resList.sort(Comparator.comparingInt(PortUse::getPort));
        return resList;
    }

    /**
     * print
     *
     * @param portUseList 端口占用信息list
     */
    private static void print(List<PortUse> portUseList) {
        TableDataPrinter printer = new TableDataPrinter(false, false);
        printer.print(PortUse.getTitle(), portUseList);
    }
}
